# 导入必要的库和模块
import torch, traceback  # PyTorch深度学习框架和错误追踪模块
import numpy as np  # 数值计算库
from Utils.tensor_ops import _2tensor  # 张量转换工具函数
from config import GlobalConfig as cfg  # 全局配置
from Utils.gpu_share import GpuShareUnit  # GPU共享管理单元
from Common.trajectory_sampler import TrajectoryDataSampler  # 轨迹数据采样器


class Trainer():
    """
    强化学习训练器类
    负责模型的核心训练逻辑，包括：
    - GPU资源管理
    - 内存安全检查
    - 轨迹数据处理
    - 损失计算和反向传播
    - 训练指标记录
    """
    def __init__(self, model, train_config):
        """
        初始化训练器

        参数:
            model: 要训练的策略网络模型
            train_config: 训练配置对象，包含各种超参数
        """
        # 核心训练组件
        self.policy = model  # 策略网络模型
        self.train_epoch = train_config.train_epoch  # 训练轮数
        self.use_avail_act = train_config.action_filter_enabled  # 是否使用可用动作过滤
        self.max_grad_norm = train_config.max_grad_norm  # 梯度裁剪阈值
        self.memory_safety_check = train_config.memory_safety_check  # 内存安全检查标志
        self.lr = train_config.learning_rate  # 学习率
        self.gamma = train_config.gamma  # 折扣因子

        # 训练状态和统计信息
        self.train_update_cnt = 0  # 训练更新次数计数器
        self.trivial_dict = {}  # 训练指标字典

        # GPU资源管理
        self.gpu_share_unit = GpuShareUnit(cfg.device, gpu_party=cfg.gpu_party)

    def train_model_on_trajectories(self, traj_pool, task):
        """
        在轨迹数据上训练模型
        包含GPU内存管理和错误处理机制。

        参数:
            traj_pool: 轨迹池，包含训练数据
            task: 任务标识符
        """
        while True:
            try:
                # 在GPU资源共享单元的保护下执行训练
                with self.gpu_share_unit:
                    self.learn_from_trajectories(traj_pool, task)
                break  # 成功完成训练，退出循环
            except RuntimeError as err:
                # 处理GPU内存不足错误
                print(traceback.format_exc())  # 打印错误堆栈信息
                if self.memory_safety_check:
                    # 内存安全检查模式下，动态调整采样大小
                    # 如果当前采样大小已为负值，则回退到前一个设置
                    if TrajectoryDataSampler.MaxSampleNum[-1] < 0:
                        TrajectoryDataSampler.MaxSampleNum.pop(-1)
                    assert TrajectoryDataSampler.MaxSampleNum[-1] > 0
                    # 标记需要减小采样大小
                    TrajectoryDataSampler.MaxSampleNum[-1] = -1
                    print('Insufficient gpu memory, using previous sample size !')
                else:
                    # 非内存安全检查模式下，直接抛出错误
                    assert False
            # 清理GPU缓存
            torch.cuda.empty_cache()

    def learn_from_trajectories(self, traj_pool, task):
        """
        从轨迹数据中学习
        执行多轮训练，每轮都从轨迹池中采样数据进行训练。

        参数:
            traj_pool: 轨迹池对象
            task: 任务标识符

        返回:
            int: 训练更新次数
        """
        # 创建轨迹数据采样器
        sampler = TrajectoryDataSampler(n_div=1, traj_pool=traj_pool, flag=task,
                                        memory_safety_check=self.memory_safety_check)

        # 执行多轮训练
        for e in range(self.train_epoch):
            # 重置采样器并获取迭代器
            sample_iter = sampler.reset_and_get_iter()
            # 获取轨迹采样数据
            train_sample = next(sample_iter)
            # 执行一轮训练
            loss_final, loss_dict = self.run_training(task, train_sample)
            # 对损失进行缩放调整
            loss_final = loss_final * 0.5

            # 在第一轮训练时打印GPU内存使用情况
            if e == 0:
                print('[PPO.py] Memory Allocated %.2f GB' % (torch.cuda.memory_allocated() / 1073741824))

            # 记录训练指标
            self.record_training_metrics(dictionary=loss_dict)
            loss_dict = None  # 释放内存

        # 完成所有轮次的更新

        # 打印训练摘要
        self.print_training_summary()
        # 更新训练计数器
        self.train_update_cnt += 1
        return self.train_update_cnt

    def record_training_metrics(self, dictionary):
        """
        记录训练指标
        将各种训练损失和性能指标记录到字典中，用于后续统计分析。

        参数:
            dictionary (dict): 包含各种训练指标的字典
        """
        for key in dictionary:
            # 如果是新指标，初始化空列表
            if key not in self.trivial_dict:
                self.trivial_dict[key] = []

            # 提取数值（如果是张量则转换为标量）
            item = dictionary[key].item() if hasattr(dictionary[key], 'item') else dictionary[key]
            # 添加到指标列表中
            self.trivial_dict[key].append(item)

    def print_training_summary(self, output=True):
        """
        打印训练摘要信息
        计算并显示各项训练指标的平均值，提供训练过程的整体概览。

        参数:
            output (bool): 是否输出打印信息
        """
        # 将所有指标列表转换为numpy数组
        for key in self.trivial_dict:
            self.trivial_dict[key] = np.array(self.trivial_dict[key])

        # 准备打印缓冲区
        print_buf = ['[Trainer.py] ']

        # 计算每个指标的平均值并格式化输出
        for key in self.trivial_dict:
            self.trivial_dict[key] = self.trivial_dict[key].mean()
            print_buf.append(' %s:%.3f, '%(key, self.trivial_dict[key]))

        # 输出摘要信息
        if output:
            print(''.join(print_buf))

        # 清空指标字典，为下一轮训练做准备
        self.trivial_dict = {}

    def run_training(self, flag, sample):
        """
        执行具体的训练步骤
        处理采样数据，计算损失，执行反向传播。

        参数:
            flag (str): 标识当前是训练还是验证
            sample (dict): 包含训练样本的字典，包含状态、动作、奖励等

        返回:
            tuple: (loss, loss_dict) - 损失值和损失字典
        """
        # 将采样数据转换为PyTorch张量
        state = _2tensor(sample['state'])      # 状态张量
        reward = _2tensor(sample['reward'])    # 奖励张量
        action = _2tensor(sample['action'])    # 动作张量
        next_state = _2tensor(sample['next_state'])  # 下一状态张量
        done = _2tensor(sample['done'])        # 终止标志张量
        avail_act = _2tensor(sample['avail_act']) if 'avail_act' in sample else None  # 可用动作掩码

        # 确保是训练模式
        assert flag == 'train'

        # 使用策略网络评估动作并计算损失
        # 这里调用的是具体的强化学习算法（如PPO、A2C等）的损失计算方法
        loss = self.policy.evaluate_actions(state_batch=state,
                                            action_batch=action,
                                            avail_act_batch=avail_act,
                                            next_state_batch=next_state,
                                            reward_batch=reward,
                                            done_batch=done,
                                            gamma=self.gamma,
                                            max_grad_norm=self.max_grad_norm)

        # 构建损失字典，用于记录和监控
        loss_dict = {"value_l1loss": loss}
        return loss, loss_dict
